package com.niodata.md.rdb.crawler;

import com.niodata.md.rdb.model.Table;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.List;

/**
 * get table storage size.
 */
public class TableStorageSizeRetrieval extends AbstractRetrieval {
  @Override
  public List<MetaDataRetrieval> retrieveMetaData() throws Exception {
    Table table = this.getCatalog().getTable(this.getParentPath());
    Connection connection = this.getDataSourceInfo().getConnection(this.getCatalog().getName());
    try {
      Statement statement = connection.createStatement();
      String sql = this.getTableSizeSql(table.getDatabase(), table.getSchema(), table.getName());
      ResultSet rs = statement.executeQuery(sql);
      rs.next();
      table.setTableSize(rs.getLong(3));
    } finally {
      connection.close();
    }
    return List.of();
  }

  private String getTableSizeSql(String database, String schema, String tableName) {
    switch (this.getCatalog().getDbInfo().getDatabaseServerType()) {
      case "mysql", "doris":
        return "select table_schema, table_name, data_length + index_length as size "
            + "from information_schema.tables where table_schema = '" + database + "' and table_name = '"
            + tableName + "'";
      case "postgresql":
        return "select schemaname, relname, pg_total_relation_size(relid) as size from "
            + "pg_stat_user_tables where schemaname = '" + schema + "' and relname = '" + tableName + "'";
      default:
        return "";
    }
  }
}
